# encoding: utf-8
# @File  : node_manager.py
# @Author: shaoyun
# @Date  :  2025/05/11
import asyncio
from loguru import logger
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from functools import wraps
import time

from sqlalchemy.orm import Session
from app.db.redis_manager import RedisClient
from app.models.node import Node, NodeStatus
from app.core.config import settings
from app.utils.metrics import measure_time

class NodeManagerException(Exception):
    """节点管理器异常"""
    pass

def handle_db_errors(func):
    """数据库操作错误处理装饰器"""

    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except Exception as e:
            logger.error(f"Database operation error in {func.__name__}: {str(e)}")
            raise NodeManagerException(f"Database operation failed: {str(e)}")

    return wrapper

class NodeManager:
    """节点管理器"""

    def __init__(self):
        self._cache = {}
        self._last_cleanup = time.time()
        self._cache_ttl = 300  # 缓存有效期5分钟

    @staticmethod
    @handle_db_errors
    @measure_time
    async def register_node(
            db: Session,
            name: str,
            host_ip: str,  # 新增 host_ip 参数
            ip: str,
            port: int,
            is_master: bool = False,
            description: Optional[str] = None
    ) -> Node:
        """注册新节点"""
        node_id = str(uuid.uuid4())

        node = Node(
            id=node_id,
            name=name,
            host_ip=host_ip,  # 设置 host_ip
            ip=ip,
            port=port,
            is_master=is_master,
            status=NodeStatus.ONLINE if is_master else NodeStatus.OFFLINE,
            description=description
        )

        db.add(node)
        db.commit()
        db.refresh(node)

        # 将节点信息缓存到Redis
        await RedisClient.set_node_info(node_id, {
            'id': node.id,
            'name': node.name,
            'host_ip': node.host_ip,  # 新增 host_ip 到缓存
            'ip': node.ip,
            'port': node.port,
            'status': node.status.value
        })

        logger.info(f"Node {node_id} registered: {name} ({ip}:{port})")
        return node

    @staticmethod
    async def get_node(db: Session, node_id: str) -> Optional[Node]:
        """获取节点详情"""
        return db.query(Node).filter(Node.id == node_id).first()

    @staticmethod
    async def get_nodes(
            db: Session,
            status: Optional[str] = None,
            skip: int = 0,
            limit: int = 100
    ) -> List[Node]:
        """获取节点列表"""
        query = db.query(Node)

        if status:
            try:
                node_status = NodeStatus(status)
                query = query.filter(Node.status == node_status)
            except ValueError:
                logger.warning(f"Invalid node status: {status}")

        return query.offset(skip).limit(limit).all()

    @staticmethod
    async def update_node_status(db: Session, node_id: str, status: NodeStatus) -> bool:
        """更新节点状态"""
        node = db.query(Node).filter(Node.id == node_id).first()

        if not node:
            return False

        node.status = status
        node.updated_at = datetime.utcnow()

        db.commit()

        # 更新Redis缓存
        await RedisClient.set_node_info(node_id, {
            'id': node.id,
            'name': node.name,
            'host_ip': node.host_ip,  # 新增 host_ip 到缓存
            'ip': node.ip,
            'port': node.port,
            'status': node.status.value
        })
        return True

    @staticmethod
    async def update_node_workload(
            db: Session,
            node_id: str,
            cpu_usage: float,
            memory_usage: float,
            disk_usage: Optional[Dict] = None
    ) -> bool:
        """更新节点负载信息"""
        node = db.query(Node).filter(Node.id == node_id).first()

        if not node:
            return False

        node.cpu_usage = cpu_usage
        node.memory_usage = memory_usage
        if disk_usage:
            node.disk_usage = disk_usage
        node.updated_at = datetime.utcnow()

        db.commit()

        # 更新Redis中的节点状态
        status_data = {
            "status": node.status.value,
            "cpu_usage": cpu_usage,
            "memory_usage": memory_usage,
            "disk_usage": disk_usage,
            "updated_at": int(datetime.utcnow().timestamp())
        }
        await RedisClient.update_node_status(node_id, status_data)

        # 检查负载阈值
        if cpu_usage > settings.CPU_THRESHOLD or memory_usage > settings.MEMORY_THRESHOLD:
            logger.warning(f"Node {node_id} high load detected: CPU {cpu_usage}%, Memory {memory_usage}%")

        return True

    @staticmethod
    async def update_node_heartbeat(db: Session, node_id: str) -> bool:
        """更新节点心跳"""
        node = db.query(Node).filter(Node.id == node_id).first()

        if not node:
            return False

        node.last_heartbeat = datetime.utcnow()
        node.status = NodeStatus.ONLINE
        db.commit()

        # 更新Redis中的心跳信息
        await RedisClient.set_heartbeat(node_id)

        return True

    @staticmethod
    async def disable_node(db: Session, node_id: str) -> bool:
        """禁用节点"""
        return await NodeManager.update_node_status(db, node_id, NodeStatus.DISABLED)

    @staticmethod
    async def enable_node(db: Session, node_id: str) -> bool:
        """启用节点"""
        return await NodeManager.update_node_status(db, node_id, NodeStatus.ONLINE)

    @staticmethod
    async def check_nodes_health(db: Session) -> Dict[str, int]:
        """检查所有节点的健康状态"""
        now = datetime.utcnow()
        timeout_threshold = now - timedelta(seconds=settings.NODE_HEARTBEAT_TIMEOUT)

        # 获取所有节点
        nodes = db.query(Node).all()

        total = len(nodes)
        online = 0
        offline = 0

        for node in nodes:
            if node.status == NodeStatus.DISABLED:
                continue

            # 检查心跳超时
            if node.last_heartbeat and node.last_heartbeat < timeout_threshold:
                if node.status != NodeStatus.OFFLINE:
                    node.status = NodeStatus.OFFLINE
                    logger.warning(f"Node {node.id} ({node.name}) is offline due to heartbeat timeout")
                offline += 1
            elif node.status == NodeStatus.ONLINE:
                online += 1

        db.commit()

        return {
            "total": total,
            "online": online,
            "offline": offline
        }

    @classmethod
    async def run_health_check_loop(cls):
        """运行健康检查循环"""
        from app.db.mysql import SessionLocal

        while True:
            try:
                db = SessionLocal()
                result = await cls.check_nodes_health(db)

                if result["offline"] > 0:
                    logger.warning(f"Detected {result['offline']} node(s) offline out of {result['total']}")

                # 发送健康检查指标到监控系统
                await cls._send_health_metrics(result)

                db.close()
                await asyncio.sleep(settings.HEALTH_CHECK_INTERVAL)
            except Exception as e:
                logger.error(f"Error in health check loop: {e}")
                await asyncio.sleep(10)

    @staticmethod
    async def _send_health_metrics(result: Dict[str, int]):
        """发送健康检查指标"""
        try:
            # 这里可以添加指标收集逻辑
            logger.debug(f"Health metrics: {result}")
        except Exception as e:
            logger.error(f"Error sending health metrics: {e}")